######################Simulations for posterior predictive p-values#############
##############################estimating functions##############################

###########################estimate propensity scores###########################
cal.ps <- function(t,xa.1,params = params)
{
  est.ps <- glm(t~xa.1+0,family = binomial(link = "logit"))$fitted.values
  return(est.ps)
}

###############################outcome regression###############################
cal.y <- function(y,x,t,params = params)
{
  lm(y~x + 0,weights = t)$fitted.values
}

##############################summary statistics################################
#############################The IPW estimator##################################
cal.tau.ipw <- function(t,xa.1,y,ps,params = params,s.ipw = NA,center = 0) 
{
#####the formula for IPW estimator#####
  tau <- sum(t*y/ps)/sum(t/ps)-sum((1-t)/(1-ps)*y)/sum((1-t)/(1-ps)) - center
  if(params$normalize){
#####if the bootstrap sd estimator is not given, it uses asymptotic formula to estimate sd#####
    if(is.na(s.ipw)){
      y1.fitted <- sum(t*y/ps)/sum(t/ps)
      y0.fitted <- sum((1-t)/(1-ps)*y)/sum((1-t)/(1-ps))
      
      Ia <- t(xa.1)%*%(c(ps*(1-ps))*xa.1)/N
      Ha <- -c((t*(y - y1.fitted)/ps^2 +
                  (1-t)*(y - y0.fitted)/(1-ps)^2) * 
                 ps * (1-ps)) * xa.1
      Ha <- t(solve(Ia,apply(Ha,2,mean)))
      psi <- t*(y-y1.fitted)/ps - (1-t)/(1-ps)*(y-y0.fitted) + c(Ha%*%t(xa.1))*(t-ps)
      s.ipw <- sqrt(var(psi)/N)
      tau.normalized = tau/s.ipw
      return(tau.normalized)
    }
    else{
      tau.normalized = tau/s.ipw
      return(tau.normalized)
    }
  }
  return(tau)
}
##########################The Regression imputation estimator###################
cal.tau.reg <- function(t,xb1.1,xb0.1,y,y1.fitted,y0.fitted,params = params,s.reg = NA, center = 0)
{
##########################the formula for reg estimator#########################
  mu.now.1 <- mean(y1.fitted)
  mu.now.0 <- mean(y0.fitted)
  tau = mu.now.1 - mu.now.0 - center
  if(params$normalize){
    if(is.na(s.reg)){
#####if the bootstrap sd estimator is not given, it uses asymptotic formula to estimate sd#####
      I1 <- t(xb1.1)%*%(c(t)*xb1.1)/N
      I0 <- t(xb0.1)%*%(c(1-t)*xb0.1)/N
      H1 <- t(solve(I1,apply(xb1.1,2,mean)))
      H0 <- t(solve(I0,apply(xb0.1,2,mean)))
      
      psi <- y1.fitted - y0.fitted +
        c(H1%*%t(xb1.1))*t*(y-y1.fitted) -
        c(H0%*%t(xb0.1))*(1-t)*(y-y0.fitted)
      s.reg <- sqrt(var(psi)/N)
      tau.normalized = tau/s.reg
      return(tau.normalized)
    }
    else{
      return(tau/s.reg)
    }
  }
  return(tau)
}
############################The doubly robust estimator#########################
cal.tau.dr <- function(t,xa.1,xb1.1,xb0.1,y,y1.fitted,y0.fitted,ps,dta,params = params,s.dr = NA,center = 0)
{
############################the formula for dr estimator########################
  tau <- mean(t/ps*(y-y1.fitted) - 
                (1-t)/(1-ps)*(y-y0.fitted) +
                y1.fitted - y0.fitted) - center
  if(params$normalize){
    if(is.na(s.dr)){
#####if the bootstrap sd estimator is not given, it uses asymptotic formula to estimate sd#####
      Ia <- t(xa.1)%*%(c(ps*(1-ps))*xa.1)/N
      I1 <- t(xb1.1)%*%(c(t)*xb1.1)/N
      I0 <- t(xb0.1)%*%(c(1-t)*xb0.1)/N
      
      Ha <- -c((t*(y - y1.fitted)/ps^2 +
                  (1-t)*(y - y0.fitted)/(1-ps)^2) * 
                 ps * (1-ps)) * xa.1
      Ha <- t(solve(Ia,apply(Ha,2,mean)))
      H1 <- c(1-t/ps)*xb1.1
      H1 <- t(solve(I1,apply(H1,2,mean)))
      H0 <- c(1-(1-t)/(1-ps))*xb0.1
      H0 <- t(solve(I0,apply(H0,2,mean)))
      
      psi <- t/ps*(y-y1.fitted) -
        (1-t)/(1-ps)*(y-y0.fitted) +
        y1.fitted - y0.fitted +
        c(Ha%*%t(xa.1))*(t-ps) +
        c(H1%*%t(xb1.1))*t*(y-y1.fitted) -
        c(H0%*%t(xb0.1))*(1-t)*(y-y0.fitted)
      
      s.dr <- sqrt(var(psi)/N)
      tau.normalized = tau/s.dr
      return(tau.normalized)
    }
    else{
      return(tau/s.dr)
    }
  }
  return(tau)
}

#####Return all statistics for all model specification##########################
#####The following three functions are just simplifications of this function####
cal.tau.allstat.allmodel <- function(dta,params = params,s = NA,center = rep(0,10)){
  t.t <- c(dta$t.t)
  t.f <- c(dta$t.f)
  xat <- dta$xat
  x1t <- dta$x1t
  x0t <- dta$x0t
  xaf <- dta$xaf
  x1f <- dta$x1f
  x0f <- dta$x0f
  y <- dta$y
#########the propensity score estimated by true model and the false one#########
  ps.t <- cal.ps(t.t,xat,params)
  ps.f <- cal.ps(t.f,xaf,params)
#########the outcome regression for four kinds of model specifications##########
  y1.fitted.tt <- lm(y~x1t + 0,weights = t.t)$fitted.values
  y0.fitted.tt <- lm(y~x0t + 0,weights = 1 - t.t)$fitted.values
  y1.fitted.ff <- lm(y~x1f + 0,weights = t.f)$fitted.values
  y0.fitted.ff <- lm(y~x0f + 0,weights = 1 - t.f)$fitted.values
  y1.fitted.tf <- lm(y~x1t + 0,weights = t.f)$fitted.values
  y0.fitted.tf <- lm(y~x0t + 0,weights = 1 - t.f)$fitted.values
  y1.fitted.ft <- lm(y~x1f + 0,weights = t.t)$fitted.values
  y0.fitted.ft <- lm(y~x0f + 0,weights = 1 - t.t)$fitted.values
################################Test statistics#################################
  stat.ipw <- (c(cal.tau.ipw(t.t,xat,y,ps.t,params = params,s.ipw = s[1],center = center[1]),
                 cal.tau.ipw(t.f,xaf,y,ps.f,params = params,s.ipw = s[2],center = center[2])))
  stat.reg <- c(cal.tau.reg(t.t,x1t,x0t,y,y1.fitted.tt,y0.fitted.tt,params = params,s.reg = s[3], center = center[3]),
                cal.tau.reg(t.t,x1f,x0f,y,y1.fitted.ft,y0.fitted.ft,params = params,s.reg = s[4], center = center[4]),
                cal.tau.reg(t.f,x1t,x0t,y,y1.fitted.tf,y0.fitted.tf,params = params,s.reg = s[5], center = center[5]),
                cal.tau.reg(t.f,x1f,x0f,y,y1.fitted.ff,y0.fitted.ff,params = params,s.reg = s[6], center = center[6]))
  stat.dr <- c(cal.tau.dr(t.t,xat,x1t,x0t,y,y1.fitted.tt,y0.fitted.tt,ps.t,dta,params = params,s.dr = s[7],center = center[7]),
               cal.tau.dr(t.t,xa.1 = xat,xb1.1 = x1f,xb0.1 = x0f,y,y1.fitted.ft,y0.fitted.ft,ps.t,dta,params = params,s.dr = s[8],center = center[8]),
               cal.tau.dr(t.f,xa.1 = xaf,xb1.1 = x1t,xb0.1 = x0t,y,y1.fitted.tf,y0.fitted.tf,ps.f,dta,params = params,s.dr = s[9],center = center[9]),
               cal.tau.dr(t.f,xa.1 = xaf,xb1.1 = x1f,xb0.1 = x0f,y,y1.fitted.ff,y0.fitted.ff,ps.f,dta,params = params,s.dr = s[10],center = center[10]))
  return(c(stat.ipw,stat.reg,stat.dr))
}

#############################Return all statistics##############################
cal.tau.allstat <- function(dta,params = params,s = NA,center = rep(0,3)){
  N <<- length(dta$t.t)
  t.t <- c(dta$t.t)
  xat <- dta$xat
  x1t <- dta$x1t
  x0t <- dta$x0t
  y <- dta$y
  ps.t <- cal.ps(t.t,xat,params)
  y1.fitted.tt <- lm(y~x1t + 0,weights = t.t)$fitted.values
  y0.fitted.tt <- lm(y~x0t + 0,weights = 1 - t.t)$fitted.values
  stat.ipw <- c(cal.tau.ipw(t.t,xat,y,ps.t,params = params,s.ipw = s[1],center = center[1]))
  stat.reg <- c(cal.tau.reg(t.t,x1t,x0t,y,y1.fitted.tt,y0.fitted.tt,params = params,s.reg = s[2], center = center[2]))
  stat.dr <- c(cal.tau.dr(t.t,xat,x1t,x0t,y,y1.fitted.tt,y0.fitted.tt,ps.t,dta,params = params,s.dr = s[3],center = center[3]))
  return(c(stat.ipw,stat.reg,stat.dr))
}

#####Return regression imputation estimators under four model specifications####
cal.tau.reg.allmodel <- function(dta,params = params,s = NA,center = rep(0,4)){
  t.t <- c(dta$t.t)
  t.f <- c(dta$t.f)
  xat <- dta$xat
  x1t <- dta$x1t
  x0t <- dta$x0t
  xaf <- dta$xaf
  x1f <- dta$x1f
  x0f <- dta$x0f
  y <- dta$y
  ps.t <- cal.ps(t.t,xat,params)
  ps.f <- cal.ps(t.f,xaf,params)
  y1.fitted.tt <- lm(y~x1t + 0,weights = t.t)$fitted.values
  y0.fitted.tt <- lm(y~x0t + 0,weights = 1 - t.t)$fitted.values
  y1.fitted.ff <- lm(y~x1f + 0,weights = t.f)$fitted.values
  y0.fitted.ff <- lm(y~x0f + 0,weights = 1 - t.f)$fitted.values
  y1.fitted.tf <- lm(y~x1t + 0,weights = t.f)$fitted.values
  y0.fitted.tf <- lm(y~x0t + 0,weights = 1 - t.f)$fitted.values
  y1.fitted.ft <- lm(y~x1f + 0,weights = t.t)$fitted.values
  y0.fitted.ft <- lm(y~x0f + 0,weights = 1 - t.t)$fitted.values
  stat.reg <- c(cal.tau.reg(t.t,x1t,x0t,y,y1.fitted.tt,y0.fitted.tt,params = params,s.reg = s, center = center[1]),
                cal.tau.reg(t.t,x1f,x0f,y,y1.fitted.ft,y0.fitted.ft,params = params,s.reg = s, center = center[2]),
                cal.tau.reg(t.f,x1t,x0t,y,y1.fitted.tf,y0.fitted.tf,params = params,s.reg = s, center = center[3]),
                cal.tau.reg(t.f,x1f,x0f,y,y1.fitted.ff,y0.fitted.ff,params = params,s.reg = s, center = center[4]))
  return(stat.reg)
}

#####Return doubly robust estimators under four model specifications############
cal.tau.dr.allmodel <- function(dta,params = params,s = NA,center = rep(0,4)){
  t.t <- c(dta$t.t)
  t.f <- c(dta$t.f)
  xat <- dta$xat
  x1t <- dta$x1t
  x0t <- dta$x0t
  xaf <- dta$xaf
  x1f <- dta$x1f
  x0f <- dta$x0f
  y <- dta$y
  ps.t <- cal.ps(t.t,xat,params)
  ps.f <- cal.ps(t.f,xaf,params)
  y1.fitted.tt <- lm(y~x1t + 0,weights = t.t)$fitted.values
  y0.fitted.tt <- lm(y~x0t + 0,weights = 1 - t.t)$fitted.values
  y1.fitted.ff <- lm(y~x1f + 0,weights = t.f)$fitted.values
  y0.fitted.ff <- lm(y~x0f + 0,weights = 1 - t.f)$fitted.values
  y1.fitted.tf <- lm(y~x1t + 0,weights = t.f)$fitted.values
  y0.fitted.tf <- lm(y~x0t + 0,weights = 1 - t.f)$fitted.values
  y1.fitted.ft <- lm(y~x1f + 0,weights = t.t)$fitted.values
  y0.fitted.ft <- lm(y~x0f + 0,weights = 1 - t.t)$fitted.values
  stat.dr <- c(cal.tau.dr(t.t,xat,x1t,x0t,y,y1.fitted.tt,y0.fitted.tt,ps.t,dta,params = params,s.dr = s[7],center = center[7]),
               cal.tau.dr(t.t,xa.1 = xat,xb1.1 = x1f,xb0.1 = x0f,y,y1.fitted.ft,y0.fitted.ft,ps.t,dta,params = params,s.dr = s[8],center = center[8]),
               cal.tau.dr(t.f,xa.1 = xaf,xb1.1 = x1t,xb0.1 = x0t,y,y1.fitted.tf,y0.fitted.tf,ps.f,dta,params = params,s.dr = s[9],center = center[9]),
               cal.tau.dr(t.f,xa.1 = xaf,xb1.1 = x1f,xb0.1 = x0f,y,y1.fitted.ff,y0.fitted.ff,ps.f,dta,params = params,s.dr = s[10],center = center[10]))
  return(stat.dr)
}

##################################THE MAIN FUNCTION#############################
#when s = "stu" we might studentize tau with asymptotic formula or not studentize it at all,
#depending whether params$normalize == T
#when s = "boot" we will use bootstrap sd estimator to studentize the test statistics
cal.tau <- function(dta,params = params,stu = "asym", c = 0)
{
########Use the value of params$statistic to determine which function above to use#######
  tau.use <- switch(params$statistic,"all"=cal.tau.allstat.allmodel,"reg"=cal.tau.reg.allmodel,"real"=cal.tau.allstat)
########The center of test statistics, we didn't use this function for this simulation#######
  c <- switch(params$statistic,"all"=rep(0,10),"reg"=rep(0,4),"real"=rep(0,3))
  if(stu == "asym")
    tryCatch(tau.use(dta,params,NA,c),error = function(e){NA})
  else if(stu == "boot"){
    boot.sd <- apply(cal.tau.boot(dta,params),1,sd,na.rm = T)
########occasionally there will be singular information matrix, so we use trycatch to return NA if such problem happens#######
    tryCatch(tau.use(dta,params,s = boot.sd,c),error = function(e){NA})
  }
}

################################The Bootstrap function##########################
cal.tau.boot <- function(dta,params,c = F)
{
  params$normalize <- F
  center <- cal.tau(dta,params)
  if(c == T){
    params$normalize <- T
  }
  else{
    center <- rep(0,length(center))
  }
  bootreplica <- function(dta.copy,params)
  {
    IndexBoot <- sample(1:N,N,replace = T)
    dta$y <- dta.copy$y[IndexBoot]
    dta$xat <- dta.copy$xat[IndexBoot,]
    dta$xaf <- dta.copy$xaf[IndexBoot,]
    dta$x1t <- dta.copy$x1t[IndexBoot,]
    dta$x1f <- dta.copy$x1f[IndexBoot,]
    dta$x0t <- dta.copy$x0t[IndexBoot,]
    dta$x0f <- dta.copy$x0f[IndexBoot,]
    dta$t.t <- dta.copy$t.t[IndexBoot]
    dta$t.f <- dta.copy$t.f[IndexBoot]
    cal.tau(dta,params,c = center)
  }
  replicate(params$bootm,bootreplica(dta,params))
}


